# coding=utf-8
#*****************************************************************************
#       Copyright (C) 2008 Sébastien Labbé <slabqc@gmail.com>
#                          Arnaud bergeron <abergeron@gmail.com>
#
#  Distributed under the terms of the GNU General Public License version 2 (GPLv2)
#
#  The full text of the GPLv2 is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from sage.structure.sage_object import SageObject
from sage.words.utils import *

class SetElement(object):
    __slots__ = ['obj', 'parent', 'rank']
    def __init__(self, obj, parent=None, rank=0):
        r"""
        Helper class for DisjointSets
        
        TESTS:
            sage: from sage.words.disjoint_set import SetElement
            sage: s = SetElement("ab")
            sage: s.obj == 'ab'
            True
            sage: s.rank
            0
            sage: s.parent is s
            True
            sage: s.other_attribute = 'abc'
            Traceback (most recent call last):
            ...
            AttributeError: 'SetElement' object has no attribute 'other_attribute'
            sage: s2 = loads(dumps(s))
            sage: s2.obj
            'ab'
            sage: s2.rank
            0
            sage: s2.parent is s2
            True
        """
        if parent is None:
            parent = self
        self.obj = obj
        self.parent = parent
        self.rank = rank

def f(x):
    r"""
    identity function
    
    TESTS:
        sage: from sage.words.disjoint_set import f
        sage: f('ab')
        'ab'
    """
    return x

class DisjointSet(SageObject):
    def __init__(self, iter_, key=f):
        r"""
        Construction de la partition triviale ou chaque element est 
        une classe d'equivalence.

        EXAMPLES:        
            sage: DisjointSet('abcde')  
            {{'a'}, {'b'}, {'c'}, {'d'}, {'e'}}
            sage: DisjointSet(range(6))
            {{0}, {1}, {2}, {3}, {4}, {5}}
            sage: DisjointSet(['yi',45,'cheval'])
            {{'cheval'}, {'yi'}, {45}}
            sage: DisjointSet(set([0, 1, 2, 3, 4]))
            {{0}, {1}, {2}, {3}, {4}}
        
        TESTS:
            sage: s = DisjointSet('abcde')
            sage: s == loads(dumps(s))
            True
        """
        self._elems = {}
        self._key = key
        
        for element in iter_:
            self._elems[element] = SetElement(element)
    
    def __eq__(self, other):
        r"""
        Equality relation for DisjointSets.
        
        TESTS:
            sage: DisjointSet('abcde') == DisjointSet('abcde')
            True
            sage: DisjointSet('abcde') == DisjointSet('abcd')
            False
        """
        if not isinstance(other, DisjointSet):
            return NotImplemented
        for e in self._elems.iterkeys():
            if e not in other or \
               other.representant_of(e) != self.representant_of(e):
                return False
        return True
        
    def __ne__(self, other):
        r"""
        Inequality relation for DisjointSets.
        
        TESTS:
            sage: DisjointSet('abcde') != DisjointSet('abcde')
            False
            sage: DisjointSet('abcde') != DisjointSet('abcd')
            True
        """
        if not isinstance(other, DisjointSet):
            return NotImplemented
        return not self.__eq__(other)
            
    def _repr_(self):
        r"""
        Retourne la partition en un unique str.
        
        EXAMPLES:
            sage: e = DisjointSet(range(5))
            sage: e.join(2,4); e._repr_()
            '{{0}, {1}, {2, 4}, {3}}'
            sage: e = DisjointSet(range(5))
            sage: e.join(4,2); e._repr_()
            '{{0}, {1}, {2, 4}, {3}}'
        """
        s = dict()
        for e in self._elems.itervalues():
            o = self._find(e.parent).obj
            if o not in s:
                s[o] = []
            s[o].append(e.obj)
        res = []
        for l in s.itervalues():
            l.sort(key=self._key)
            res.append('{%s}'% ', '.join(imap(repr, l)))
        res.sort()
        return '{%s}'% ', '.join(res)
    
    def representants(self):
        r"""
        Return the list of all the actual representants of the partition.
        
        EXAMPLES:
            sage: e = DisjointSet(range(5))
            sage: e
            {{0}, {1}, {2}, {3}, {4}}
            sage: e.join(2,3)
            sage: e.join(3,4)
            sage: e    
            {{0}, {1}, {2, 3, 4}} 
            sage: list(e.representants())
            [0, 1, 2]
        """
        for e in self._elems.itervalues():
            if e.parent is e:
                yield e.obj
        
    def representant_of(self, a):
        r"""
        Retourne le representant de a.
        Return None if a is not in the partition.
        
        EXAMPLES:
            sage: e = DisjointSet(range(5))
            sage: e.join(4,2)
            sage: e
            {{0}, {1}, {2, 4}, {3}}
            sage: e.representant_of(2)
            2
            sage: e.representant_of(4)
            2
            sage: e.join(1,3)
            sage: e
            {{0}, {1, 3}, {2, 4}}
            sage: e.representant_of(1)
            1
            sage: e.representant_of(3)
            1
            sage: e.join(3,2)
            sage: e
            {{0}, {1, 2, 3, 4}}
            sage: [e.representant_of(i) for i in range(5)]
            [0, 1, 1, 1, 1]
            sage: e.representant_of(5)
            Traceback (most recent call last):
            ...
            ValueError: 5 not in set
        """
        if a not in self._elems:
            raise ValueError, '%r not in set'%(a,)
        return self._find(self._elems[a]).obj
    
    def _find(self, e):
        r"""
        Internal find using SetElements
        
        TESTS:
            sage: s = DisjointSet(range(5))
            sage: s.join(2, 3)
            sage: s.join(1, 4)
            sage: s.join(3, 0)
            sage: s.add(6)
            sage: s.add(5)
            sage: s.join(6, 1)
            sage: s._find(s._elems[2]).obj
            0
            sage: s._find(s._elems[6]).obj
            1
            sage: s._find(s._elems[5]).obj
            5
            sage: s._find(s._elems[3]).obj
            0
        """
        if e.parent is e:
            return e
        else:
            e.parent = self._find(e.parent)
            return e.parent
        
    def __contains__(self, a):
        r"""
        Retourne VRAI si l'element a est dans la partition
        
        EXAMPLES:
            sage: e = DisjointSet(range(5))
            sage: e
            {{0}, {1}, {2}, {3}, {4}}
            sage: 4 in e                         
            True
            sage: 5 in e
            False
        """
        return a in self._elems
            
    def add(self, a):
        r"""
        Ajoute l'element a dans la partition dans une nouvelle classe. 
        Ne fait rien si l'element est deja dans la partition.
        
        EXAMPLES:
            sage: e = DisjointSet(range(5)); e
            {{0}, {1}, {2}, {3}, {4}}
            sage: e.add(3); e
            {{0}, {1}, {2}, {3}, {4}}
            sage: e.add(18); e
            {{0}, {18}, {1}, {2}, {3}, {4}}
            sage: e.add('a'); e
            {{'a'}, {0}, {18}, {1}, {2}, {3}, {4}}
            
        TESTS:
            sage: e = DisjointSet(range(5)); e
            {{0}, {1}, {2}, {3}, {4}}
            sage: e.join(2, 3); e
            {{0}, {1}, {2, 3}, {4}}
            sage: e.add(3); e
            {{0}, {1}, {2, 3}, {4}} 
        """
        if a not in self._elems:
            self._elems[a] = SetElement(a)
            
    def join(self, a, b):
        r"""
        Joindre les classes des elements a et b
        
        EXAMPLES:
            sage: e = DisjointSet('abcde')
            sage: e
            {{'a'}, {'b'}, {'c'}, {'d'}, {'e'}}
            sage: e.join('a','b')
            sage: e
            {{'a', 'b'}, {'c'}, {'d'}, {'e'}}
            sage: e.join('c','e')
            sage: e                 
            {{'a', 'b'}, {'c', 'e'}, {'d'}}
            sage: e.join('b','e')
            sage: e
            {{'a', 'b', 'c', 'e'}, {'d'}}
        """
        repA = self._find(self._elems[a])
        repB = self._find(self._elems[b])
        
        if self._key(repA.obj) < self._key(repB.obj):
            repB.parent = repA
        else:
            repA.parent = repB
    
    def nb_classes(self):
        r"""
        Retourne le nombre de classes actuelles.
        
        EXAMPLES:
            sage: e = DisjointSet(range(5))
            sage: e.nb_classes()
            5
            sage: e.join(2,3)
            sage: e
            {{0}, {1}, {2, 3}, {4}}
            sage: e.join(0,3)
            sage: e
            {{0, 2, 3}, {1}, {4}}
            sage: e.nb_classes()
            3
        """
        return len_it(self.representants())
